# -*- coding: utf-8 -*-

import logging
import string
import sys
import time

import math

reload(sys)
sys.setdefaultencoding('utf8')

FREEFALL_ACCEL = 9.80665 * 1000.
SCALE = 0.0039 * FREEFALL_ACCEL  # 3.9mg/LSB * Earth gravity in mm/s**2


def info(messgae):
    logging.info('[AQGL][INFO]:{}'.format(messgae))


class AxesMappingHelper:
    def __init__(self, get_gravity_func):

        self.get_gravity = get_gravity_func
        self.x_calibrate_vector = None
        self.y_calibrate_vector = None
        self.origin_vector = None

    def start(self, sec):
        self.origin_vector, _ = self.get_gravity(sec, enable_mapping=False)

    def test_x(self, sec):
        self.x_calibrate_vector, _ = self.get_gravity(sec, enable_mapping=False)

    def test_y(self, sec):
        self.y_calibrate_vector, _ = self.get_gravity(sec, enable_mapping=False)

    def clear(self):
        self.origin_vector = None
        self.x_calibrate_vector = None
        self.y_calibrate_vector = None

    def cal(self, send_message, save_config):
        origin_vector = [-1] + [x for x in self.origin_vector]
        z_pos = 0
        z_dir = ''
        for _z_pos, value in enumerate(origin_vector):
            if abs(value) > 8500 and abs(value) > origin_vector[z_pos]:
                z_pos = _z_pos
                z_dir = '-' if value > 0 else ''
        z_pos -= 1
        if z_pos < 0:
            self.clear()
            send_message('ADXL345 初始姿态倾角过大!')
            return
        x_pos, y_pos = -1, -1
        x_dir, y_dir = '', ''
        delta = -1
        for _pos, value in enumerate(self.x_calibrate_vector):
            if _pos == z_pos:
                continue
            _delta = value - self.origin_vector[_pos]
            if abs(_delta) > delta:
                delta = abs(_delta)
                x_pos = _pos
                x_dir = '' if _delta > 0 else '-'

        delta = -1
        for _pos, value in enumerate(self.y_calibrate_vector):
            if _pos == z_pos:
                continue
            _delta = value - self.origin_vector[_pos]
            if abs(_delta) > delta:
                delta = abs(_delta)
                y_pos = _pos
                y_dir = '' if _delta > 0 else '-'

        axes = set(range(3))
        for axis in [x_pos, y_pos, z_pos]:
            try:
                axes.remove(axis)
            except:
                pass
        if axes:
            self.clear()
            send_message('初始姿态倾角数据异常，请尝试在测量时以更大倾角倾斜打印机。')
            return

        axes_mapping = {0: 'x', 1: 'y', 2: 'z'}
        new_mapping = [axes_mapping[axis] for axis in [x_pos, y_pos, z_pos]]

        axes_mapping_config = '{}{},{}{},{}{}'.format(x_dir, new_mapping[0], y_dir, new_mapping[1], z_dir,
                                                      new_mapping[2])
        save_config('adxl345', 'axes_map', axes_mapping_config)
        send_message('axes_map:{}\n输入 SAVE_CONFIG 保存坐标轴映射.'.format(axes_mapping_config))


class GravityLevel:
    def __init__(self, config):

        self.config = config
        self.printer = config.get_printer()
        self.gcode = self.printer.lookup_object('gcode')
        self.gcode.register_command('GET_GRAVITY', self.cmd_GET_GRAVIRY)
        self.gcode.register_command('RECORD_GRAVITY', self.cmd_RECORD_GRAVITY)

        self.gcode.register_command('AXES_MAPPING_CALIBRATE', self.cmd_AXES_MAPPING_CALIBRATE)
        self.gcode.register_command('GQGL', self.cmd_GQGL)
        self.gcode.register_command('ATEST', self.cmd_test)
        self.printer.register_event_handler("klippy:ready", self._handle)

        bed_size = self.config.getlist('bed_size', count=2)
        self.bed_size = [float(size.strip()) for size in bed_size]
        self.safe_z = config.getfloat('safe_z',-1)
        self.max_x,self.max_y = self.bed_size
        self.leveling_times = self.config.getint('leveling_times', 4)
        self.measure_second = self.config.getfloat('measure_second', 0.8)
        self.speed = self.config.getfloat('speed', -1)
        recorded_gravity = self.config.getlist('gravity', ('0.', '0.', '0.'), count=3)
        self.recorded_gravity = tuple([float(g.strip()) for g in recorded_gravity])
        self.axes_mapping_helper = None
        self.debug = False

    def send_message(self, message):
        message = str(message)
        self.gcode.respond_info(message)

    def move(self, pos, speed):
        self.printer.lookup_object('toolhead').manual_move(pos, speed)

    def get_position(self):
        return self.printer.lookup_object('toolhead').get_position()

    def save_config(self, secltion, option, data):
        configfile = self.printer.lookup_object('configfile')
        configfile.set(secltion, option, data)

    def cmd_test(self, gcmd):
        aclient = self.adxl345.start_internal_client()
        toolhead = self.printer.lookup_object('toolhead')
        toolhead.wait_moves()
        self.move([180,150],300)
        aclient.finish_measurements()
        values = aclient.get_samples()
        if not values:
            raise gcmd.error("No adxl345 measurements found")
        vectors = []
        for count, [_, x, y, z] in enumerate(values):
            vectors.append((x,y,z))
        for v in vectors:
            self.send_message(v)
        self.send_message(len(values))


    def cmd_AXES_MAPPING_CALIBRATE(self, gcmd):
        is_start = gcmd.get_int('START', 0)
        is_x_mapping_calibrate = gcmd.get_int('X', 0)
        is_y_mapping_calibrate = gcmd.get_int('Y', 0)
        wait = gcmd.get_int('WAIT', 5)

        if is_start:
            self.axes_mapping_helper = AxesMappingHelper(self.get_gravity_vector)
            self.send_message('开始进行 AXDL345 坐标轴映射校准。\n测量初始姿态:')
            self.axes_mapping_helper.start(1)
            self.send_message(
                '测量完成!\n请执行 "AXES_MAPPING_CALIBRATE x=1" 进行 X 方向姿态测量.')
            return
        if is_x_mapping_calibrate:
            if not self.axes_mapping_helper or not self.axes_mapping_helper.origin_vector:
                self.send_message('请先执行 "AXES_MAPPING_CALIBRATE start=1"')
                return
            self.printer.lookup_object('toolhead').dwell(1.)
            while wait > 0:
                self.send_message('{}s 后开始测量...'.format(wait))
                self.printer.lookup_object('toolhead').dwell(1.)
                wait -= 1
            self.send_message('测量 X 方向姿态:')
            self.axes_mapping_helper.test_x(1)
            self.send_message(
                '测量完成!\n请执行 "AXES_MAPPING_CALIBRATE y=1" 进行 y 方向姿态测量.')
            return
        if is_y_mapping_calibrate:
            if not self.axes_mapping_helper or not self.axes_mapping_helper.origin_vector:
                self.send_message('请先执行 AXES_MAPPING_CALIBRATE start=1')
                return
            if not self.axes_mapping_helper.x_calibrate_vector:
                self.send_message('请先进行 x 方向姿态测量')
                return
            self.printer.lookup_object('toolhead').dwell(1.)
            while wait > 0:
                self.send_message('{}s 后开始测量...'.format(wait))
                self.printer.lookup_object('toolhead').dwell(1.)
                wait -= 1
            self.send_message('正在记录 y 方向姿态..')
            self.axes_mapping_helper.test_y(1)
            self.send_message('测量完成!')
            self.axes_mapping_helper.cal(self.send_message, self.save_config)

    def _stepper_list(self, adj, speed):
        max_adj = self.qgl.max_adjust
        adj = [min(a, max_adj) for a in adj]
        adj = [max(a, -max_adj) for a in adj]
        self.qgl.z_helper.adjust_steppers(adj, speed)

    def level_x(self, angle_x, pos, times, sec=1.):
        abs_delta = []
        for i in range(times):
            g_vector, _ = self.get_gravity_vector(sec, show_info=False)
            current_angle_x = math.degrees(math.atan(g_vector[0] / g_vector[2]))
            delta_angle = current_angle_x - angle_x
            lift_left = math.tan(math.radians(delta_angle)) * self.max_x

            z2, z3 = 0., 0.
            if pos == 'behind':
                z0, z1 = lift_left, lift_left
                self.send_message('后方 X 倾角: {:.4f} 目标倾角:{:.4f} 差值:{:.4f}'.format(current_angle_x, angle_x, delta_angle))
            else:
                z0, z1 = lift_left, 0.
                self.send_message('前方 X 倾角: {:.4f} 目标倾角:{:.4f} 差值:{:.4f}'.format(current_angle_x, angle_x, delta_angle))
            speed = self.qgl.probe_helper.get_lift_speed()
            adjust = [z0, z1, z2, z3]
            self.send_message('4z调整距离:[{:.5f},{:.5f},{:.5f},{:.5f}]'.format(*adjust))
            if self.debug:
                self.send_message(abs_delta+[abs(delta_angle)])
            if abs_delta:
                if abs(delta_angle) > abs_delta[-1]:
                    if abs(delta_angle) > 1 or (len(abs_delta) > 1 and (abs_delta[-1] > abs_delta[-2] or abs(delta_angle) > abs_delta[-2])):
                        self.send_message('到达加速度精度上限')
                        break
            self._stepper_list(adjust, speed)
            abs_delta.append(abs(delta_angle))

    def level_y(self, angle_x, angle_y, times, sec=1.):
        abs_delta = []
        for i in range(times):
            g_vector, _ = self.get_gravity_vector(sec, show_info=False)
            current_angle_y = math.degrees(math.atan(g_vector[1] / g_vector[2]))
            current_angle_x = math.degrees(math.atan(g_vector[0] / g_vector[2]))
            delta_angle_y = current_angle_y - angle_y
            delta_angle_x = current_angle_x - angle_x
            lift_front = math.tan(math.radians(delta_angle_y)) * self.max_y
            lift_left = math.tan(math.radians(delta_angle_x)) * self.max_x

            zs = [lift_front + lift_left ,lift_left,0.,lift_front]
            lift_mean = sum(zs) / 4

            speed = self.qgl.probe_helper.get_lift_speed()
            adjust = [z-lift_mean for z in zs]
            self.send_message('Y 倾角: {:.4f} 目标倾角:{:.4f} 差值:{:.4f}'.format(current_angle_y, angle_y, delta_angle_y))
            self.send_message('4z调整距离:[{:.5f},{:.5f},{:.5f},{:.5f}]'.format(*adjust))

            delta_angle = abs(delta_angle_x) + abs(delta_angle_y)
            if self.debug:
                self.send_message(abs_delta+[delta_angle])

            if abs_delta:
                if delta_angle > abs_delta[-1]:
                    if delta_angle > 1 or (len(abs_delta) > 1 and (abs_delta[-1] > abs_delta[-2] or delta_angle > abs_delta[-2])):
                        self.send_message('到达加速度精度上限')
                        break
            self._stepper_list(adjust, speed)
            abs_delta.append(delta_angle)

    def cmd_GQGL(self, gcmd):
        self.debug = gcmd.get_int('DEBUG', 0)
        qgl_g_vector = self.recorded_gravity
        if all([abs(g) < 1e-5   for g in qgl_g_vector]):
            self.send_message('请先记录调平状态下的重力方向')
            return
        angle_x = math.degrees(math.atan(qgl_g_vector[0] / qgl_g_vector[2]))
        angle_y = math.degrees(math.atan(qgl_g_vector[1] / qgl_g_vector[2]))
        self.send_message('标准重力方向: ({:.0f},{:.0f},{:.0f})'.format(*qgl_g_vector))
        self.send_message('标准 x 倾角:{:.4f}  标准 y 倾角:{:.4f}'.format(angle_x, angle_y))
        # speed = self.qgl.probe_helper.get_lift_speed()
        speed = 300
        times = gcmd.get_int('TIMES', self.leveling_times)
        sec = gcmd.get_float('SEC', self.measure_second)

        if self.safe_z > 0:
            new_pos = self.get_position()[:]
            new_pos[2] = self.safe_z
            self.move(new_pos, speed)
        self.move([self.max_x / 2, self.max_y - 5], speed)
        self.level_x(angle_x, 'behind', times, sec)
        self.move([self.max_x / 2, 5], speed)
        self.level_x(angle_x, 'front', times, sec)
        self.move([self.max_x / 2, self.max_y / 2], speed)
        self.level_y(angle_x, angle_y, times, sec)

    def cmd_STEPPER_LIFT(self, gcmd):
        speed = self.qgl.probe_helper.get_lift_speed()
        max_adj = self.qgl.max_adjust
        z0 = gcmd.get_float('Z0', 0.)
        z1 = gcmd.get_float('Z1', 0.)
        z2 = gcmd.get_float('Z2', 0.)
        z3 = gcmd.get_float('Z3', 0.)
        adjust = [z0, z1, z2, z3]
        self._stepper_list(adjust, speed)

    def _handle(self):
        self.adxl345 = self.printer.lookup_objects('adxl345')[0][1]
        self.qgl = self.printer.lookup_object('quad_gantry_level')
        self.speed = self.speed if self.speed > 0 else self.qgl.probe_helper.get_lift_speed()
        # self.qgl = None

    def get_gravity_vector(self, sec=0.2, enable_mapping=True, show_info=True):
        if not enable_mapping:
            axes_map = self.adxl345.axes_map
            self.adxl345.axes_map = ((0, SCALE), (1, SCALE), (2, SCALE))
        try:
            aclient = self.adxl345.start_internal_client()
            toolhead = self.printer.lookup_object('toolhead')
            count = 0
            while sec > 1:
                toolhead.dwell(1.)
                sec -= 1
                count += 1
                if show_info:
                    self.send_message('正在测量({})...'.format(count))
            toolhead.dwell(sec)
            count += 1
            if show_info:
                self.send_message('正在测量({})...'.format(count))

            aclient.finish_measurements()
            values = aclient.get_samples()
            if not values:
                raise gcmd.error("No adxl345 measurements found")
            sum_x, sum_y, sum_z = 0., 0., 0.
            for count, [_, x, y, z] in enumerate(values):
                sum_x += x
                sum_y += y
                sum_z += z
            avg_x, avg_y, avg_z = sum_x / count, sum_y / count, sum_z / count
        except Exception as e:
            raise e
        finally:
            if not enable_mapping:
                self.adxl345.axes_map = axes_map

        return (avg_x, avg_y, avg_z), count

    def cmd_GET_GRAVIRY(self, gcmd):
        sec = gcmd.get_float('SEC', 3.)
        enable_mapping = gcmd.get_int('ENABLE_MAPPING', 1)
        [avg_x, avg_y, avg_z], _ = self.get_gravity_vector(sec, enable_mapping)
        self.send_message('x:{:.0f}, y:{:.0f}, z:{:.0f}'.format(avg_x, avg_y, avg_z))

    def cmd_RECORD_GRAVITY(self, gcmd):
        sec = gcmd.get_float('SEC', 7.)
        enable_mapping = gcmd.get_int('ENABLE_MAPPING', 1)
        [avg_x, avg_y, avg_z], times = self.get_gravity_vector(sec, enable_mapping)
        self.send_message('x:{:.0f}, y:{:.0f}, z:{:.0f}'.format(avg_x, avg_y, avg_z))
        self.save_config('gqgl', 'gravity', '{},{},{}'.format(avg_x, avg_y, avg_z))
        self.send_message("使用 SAVE_CONFIG 保存重力方向.")


def load_config(config):
    return GravityLevel(config)
